#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
修复 Parquet 文件中的 tool_call 标签格式错误

功能：
    1. 将 parquet 文件中的以下格式替换为 "</tool_call>"：
       - "</ tool _ call >" -> "</tool_call>"
       - "</Tool Call >" -> "</tool_call>"
    2. 对于第一行的 metrics map 中的 key 'AIOps-22/top_5'，将其值修改为 1
    
    直接修改原始文件：data/Test/doubao-seed-1-6-thinking-250715_20251106_000758/train.parquet
    
使用方法：
    python fix_tool_call_spacing.py
"""

import sys
import pandas as pd
import numpy as np
from pathlib import Path
from typing import Any, Dict, Tuple


def replace_in_string(text: str, old: str, new: str) -> str:
    """
    在字符串中替换指定的文本
    
    参数：
        text: 输入字符串
        old: 要替换的旧文本
        new: 替换后的新文本
        
    返回：
        替换后的字符串
    """
    return text.replace(old, new)


def replace_in_data(data: Any, old: str, new: str) -> Any:
    """
    递归地在数据结构中替换文本
    
    支持的数据类型：
        - str: 直接替换
        - dict: 递归处理所有值
        - list: 递归处理所有元素
        - numpy.ndarray: 转换为列表处理后再转换回数组
        - 其他类型: 直接返回
    
    参数：
        data: 输入数据（可以是字符串、字典、列表、numpy数组等）
        old: 要替换的旧文本
        new: 替换后的新文本
        
    返回：
        处理后的数据
    """
    if isinstance(data, str):
        return replace_in_string(data, old, new)
    elif isinstance(data, dict):
        return {k: replace_in_data(v, old, new) for k, v in data.items()}
    elif isinstance(data, (list, np.ndarray)):
        # 对于 numpy array，先转换为列表处理
        if isinstance(data, np.ndarray):
            # 转换为列表进行递归处理
            processed_list = [replace_in_data(item, old, new) for item in data.tolist()]
            # 尝试保持原有的数据类型
            try:
                return np.array(processed_list, dtype=object)
            except:
                return np.array(processed_list)
        else:
            # 普通列表
            return [replace_in_data(item, old, new) for item in data]
    else:
        return data


def update_metrics_value(data: Any, target_key: str, new_value: Any) -> Tuple[Any, bool]:
    """
    递归地在数据结构中查找并更新 metrics 字典中的指定键值
    
    支持的数据类型：
        - dict: 如果包含 'metrics' 键，则更新其值；否则递归处理所有值
        - list: 递归处理所有元素
        - numpy.ndarray: 转换为列表处理后再转换回数组
        - 其他类型: 直接返回
    
    参数：
        data: 输入数据（可以是字符串、字典、列表、numpy数组等）
        target_key: 要更新的 metrics 键名（如 'AIOps-22/top_5'）
        new_value: 新的值
        
    返回：
        (处理后的数据, 是否进行了修改)
    """
    if isinstance(data, dict):
        # 如果当前字典包含 'metrics' 键，且 metrics 是字典类型
        if 'metrics' in data and isinstance(data['metrics'], dict):
            metrics = data['metrics']
            if target_key in metrics:
                # 复制字典以避免修改原始数据
                new_data = data.copy()
                new_metrics = metrics.copy()
                new_metrics[target_key] = new_value
                new_data['metrics'] = new_metrics
                return new_data, True
            else:
                # metrics 存在但不包含目标键，继续递归处理
                modified = False
                new_data = {}
                for k, v in data.items():
                    new_v, v_modified = update_metrics_value(v, target_key, new_value)
                    new_data[k] = new_v
                    modified = modified or v_modified
                return new_data, modified
        else:
            # 递归处理所有值
            modified = False
            new_data = {}
            for k, v in data.items():
                new_v, v_modified = update_metrics_value(v, target_key, new_value)
                new_data[k] = new_v
                modified = modified or v_modified
            return new_data, modified
    elif isinstance(data, (list, np.ndarray)):
        # 对于 numpy array，先转换为列表处理
        if isinstance(data, np.ndarray):
            processed_list = []
            modified = False
            for item in data.tolist():
                new_item, item_modified = update_metrics_value(item, target_key, new_value)
                processed_list.append(new_item)
                modified = modified or item_modified
            # 尝试保持原有的数据类型
            try:
                return np.array(processed_list, dtype=object), modified
            except:
                return np.array(processed_list), modified
        else:
            # 普通列表
            processed_list = []
            modified = False
            for item in data:
                new_item, item_modified = update_metrics_value(item, target_key, new_value)
                processed_list.append(new_item)
                modified = modified or item_modified
            return processed_list, modified
    else:
        # 其他类型直接返回
        return data, False


def fix_tool_call_spacing(input_path: str, 
                          replacements: list = None) -> Dict[str, Any]:
    """
    修复 parquet 文件中的 tool_call 标签格式
    
    参数：
        input_path: 输入的 parquet 文件路径（直接修改原文件）
        replacements: 替换规则列表，每个规则为 (old_text, new_text) 元组
                      默认为：[("</ tool _ call >", "</tool_call>"), ("</Tool Call >", "</tool_call>")]
        
    返回：
        包含处理统计信息的字典：
            - rows_processed: 处理的总行数
            - rows_modified: 修改的行数
            - replacements_made: 总替换次数（每个替换规则的总和）
            - metrics_updated: 是否更新了第一行的 metrics
            - output_path: 输出文件路径
    """
    # 默认替换规则
    if replacements is None:
        replacements = [
            ("</ tool _ call >", "</tool_call>"),
            ("</Tool Call >", "</tool_call>")
        ]
    
    # 读取 parquet 文件
    df = pd.read_parquet(input_path)
    
    # 统计信息
    rows_processed = len(df)
    rows_modified = 0
    replacements_made = 0
    metrics_updated = False
    
    # 处理每一行
    for idx in df.index:
        row_modified = False
        
        # 处理每一列
        for col in df.columns:
            original_value = df.at[idx, col]
            current_value = original_value
            
            # 如果是第一行，先处理 metrics 更新
            if idx == df.index[0]:
                updated_value, was_updated = update_metrics_value(
                    current_value, 
                    'AIOps-22/top_5', 
                    1
                )
                if was_updated:
                    current_value = updated_value
                    row_modified = True
                    metrics_updated = True
            
            # 对每个替换规则进行处理
            for old_text, new_text in replacements:
                # 检查是否包含需要替换的文本（转换为字符串比较）
                value_str = str(current_value)
                if old_text in value_str:
                    # 在数据中进行递归替换
                    current_value = replace_in_data(current_value, old_text, new_text)
                    row_modified = True
                    
                    # 统计替换次数
                    replacements_made += value_str.count(old_text)
            
            # 如果有修改，更新 DataFrame
            if row_modified:
                df.at[idx, col] = current_value
        
        if row_modified:
            rows_modified += 1
    
    # 直接保存到原文件
    df.to_parquet(input_path, index=False)
    
    # 返回统计信息
    return {
        'rows_processed': rows_processed,
        'rows_modified': rows_modified,
        'replacements_made': replacements_made,
        'metrics_updated': metrics_updated,
        'output_path': input_path
    }


def main():
    """
    主函数：直接修改指定的 parquet 文件
    """
    # 固定的文件路径
    input_path = 'data/sft/doubao-seed-1-6-thinking-250715_20251106_000758/train.parquet'
    
    # 替换规则列表
    replacements = [
        ("</ tool _ call >", "</tool_call>"),
        ("</Tool Call >", "</tool_call>")
    ]
    
    # 检查输入文件是否存在
    if not Path(input_path).exists():
        print(f"错误：输入文件不存在: {input_path}", file=sys.stderr)
        sys.exit(1)
    
    try:
        # 执行替换
        print(f"正在处理文件: {input_path}")
        print("替换规则:")
        for old_text, new_text in replacements:
            print(f"  '{old_text}' -> '{new_text}'")
        
        stats = fix_tool_call_spacing(input_path, replacements)
        
        # 显示统计信息
        print(f"\n处理完成！")
        print(f"  处理的总行数: {stats['rows_processed']}")
        print(f"  修改的行数: {stats['rows_modified']}")
        print(f"  总替换次数: {stats['replacements_made']}")
        if stats['metrics_updated']:
            print(f"  第一行的 metrics['AIOps-22/top_5'] 已更新为 1")
        print(f"  输出文件: {stats['output_path']}")
        
    except Exception as e:
        print(f"错误：处理文件时发生异常: {e}", file=sys.stderr)
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == '__main__':
    main()

